{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "In this notebook, we provide an example (Diving48) on how to prepare 2D skeleton data for custom video datasets.\n",
    "You can run each step in this notebook or implement them on your own. \n",
    "\n",
    "**Step 1**: Download Videos and Annotations from the official website of Diving48: http://www.svcl.ucsd.edu/projects/resound/dataset.html."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Download Videos and Annotations\n",
    "!wget http://www.svcl.ucsd.edu/projects/resound/Diving48_V2_train.json\n",
    "!wget http://www.svcl.ucsd.edu/projects/resound/Diving48_V2_test.json\n",
    "!wget http://www.svcl.ucsd.edu/projects/resound/Diving48_rgb.tar.gz\n",
    "# Extract the videos\n",
    "!tar -xf Diving48_rgb.tar.gz\n",
    "# After that, the folder structure will looks like:\n",
    "# WorkingDirectory\n",
    "# ├── Diving48_V2_train.json\n",
    "# ├── Diving48_V2_test.json\n",
    "# `-- rgb\n",
    "#     ├-- _8Vy3dlHg2w_00000.mp4\n",
    "#     ├-- _8Vy3dlHg2w_00001.mp4\n",
    "#     ├-- _8Vy3dlHg2w_00004.mp4\n",
    "#     ├-- ..."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# After extracting videos successfully, you can now remove the tar.gz file\n",
    "!rm Diving48_rgb.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 2**: Generate video list for 2d skeleton extraction"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "from mmcv import load, dump\n",
    "from pyskl.smp import *\n",
    "train = load('Diving48_V2_train.json')\n",
    "test = load('Diving48_V2_test.json')\n",
    "tmpl = 'examples/extract_diving48_skeleton/rgb/{}.mp4'\n",
    "\n",
    "lines = [(tmpl + ' {}').format(x['vid_name'], x['label']) for x in train + test]\n",
    "mwlines(lines, 'diving48.list')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 3**: Extract 2D skeleton for Diving48 Videos"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir('../..')\n",
    "# Extract 2D skeletons of diving48 videos with 8 GPUs (it may take around 12 hours). The video_list is diving48.list  the output file is diving48.pkl\n",
    "!bash tools/dist_run.sh tools/data/custom_2d_skeleton.py 8 --video-list examples/extract_diving48_skeleton/diving48.list --out examples/extract_diving48_skeleton/diving48_annos.pkl\n",
    "# If you fail to run the previous command, please try the following one (comment the previous one and uncomment the following one, or run the following one in terminal)\n",
    "# python3 tools/data/custom_2d_skeleton.py --video-list examples/extract_diving48_skeleton/diving48.list --out examples/extract_diving48_skeleton/diving48_annos.pkl --non-dist"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "**Step 4**: Merge extracted 2D skeletons and split information to get the final annotation file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "os.chdir('examples/extract_diving48_skeleton')\n",
    "train = load('Diving48_V2_train.json')\n",
    "test = load('Diving48_V2_test.json')\n",
    "annotations = load('diving48_annos.pkl')\n",
    "split = dict()\n",
    "split['train'] = [x['vid_name'] for x in train]\n",
    "split['test'] = [x['vid_name'] for x in test]\n",
    "dump(dict(split=split, annotations=annotations), 'diving48_hrnet.pkl')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now you get the 2D skeleton annotations for diving48. We also provide the generated pickle file at https://download.openmmlab.com/mmaction/pyskl/data/diving48/diving48_hrnet.pkl. You can download this pickle file to check if it's the same as your generated pickle file. We also provide a [config file](https://github.com/kennymckormick/pyskl/blob/main/configs/posec3d/slowonly_r50_gym/joint.py) for training on diving48 and release the trained weights and accuracy number in [PoseC3D](https://github.com/kennymckormick/pyskl/blob/main/configs/posec3d/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "interpreter": {
   "hash": "8516388d98336f8f25b57d66898c29467d9fc2c465b9ee1e16d04a4702b7bff1"
  },
  "kernelspec": {
   "display_name": "Python 3.8.5 ('base')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.5"
  },
  "orig_nbformat": 4
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
